#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <math.h>

#ifndef BLOCK_M
#define BLOCK_M 64
#endif
#ifndef BLOCK_N
#define BLOCK_N 64
#endif
#ifndef BLOCK_K
#define BLOCK_K 64
#endif

static double *xmalloc(size_t count)
{
    double *p = (double *)malloc(count * sizeof(double));
    if (!p)
    {
        fprintf(stderr, "Out of memory (requested %zu doubles)\n", count);
        exit(EXIT_FAILURE);
    }
    return p;
}

static void fill_random(double *a, size_t n)
{
    for (size_t i = 0; i < n; ++i)
        a[i] = (double)rand() / (double)RAND_MAX - 0.5;
}

/* Reference naive triple loop for verification (O(M*N*K)) */
static void gemm_naive(size_t M, size_t N, size_t K, const double *A, const double *B, double *C)
{
    for (size_t i = 0; i < M; ++i)
    {
        for (size_t j = 0; j < K; ++j)
        {
            double acc = 0.0;
            for (size_t p = 0; p < N; ++p)
            {
                acc += A[i * N + p] * B[p * K + j];
            }
            C[i * K + j] = acc;
        }
    }
}
/*
A tiny 2x2 microkernel computing:
C[ir:ir+2, jr:jr+2] += A[ir:ir+2, p:p+TK] * B[p:p+TK, jr:jr+2]
Operates on row-major storage.
 */
static void microkernel_2x2(size_t N, size_t K,
                            const double *A, size_t lda,
                            const double *B, size_t ldb,
                            double *C, size_t ldc)
{
    double c00 = 0.0, c01 = 0.0, c10 = 0.0, c11 = 0.0;
    for (size_t p = 0; p < N; ++p)
    {
        double a0 = A[0 * lda + p];
        double a1 = A[1 * lda + p];
        double b0 = B[p * ldb + 0];
        double b1 = B[p * ldb + 1];
        c00 += a0 * b0;
        c01 += a0 * b1;
        c10 += a1 * b0;
        c11 += a1 * b1;
    }
    C[0 * ldc + 0] += c00;
    C[0 * ldc + 1] += c01;
    C[1 * ldc + 0] += c10;
    C[1 * ldc + 1] += c11;
}

/* Blocked GEMM: C = A * B (no alpha/beta for simplicity) */
static void gemm_blocked(size_t M, size_t N, size_t K,
                         const double *A, const double *B, double *C)
{
    // Initialize C to zero.
    memset(C, 0, sizeof(double) * M * K);
    for (size_t kk = 0; kk < N; kk += BLOCK_K)
    {
        size_t Tk = (kk + BLOCK_K <= N) ? BLOCK_K : (N - kk);
        for (size_t mm = 0; mm < M; mm += BLOCK_M)
        {
            size_t Tm = (mm + BLOCK_M <= M) ? BLOCK_M : (M - mm);
            for (size_t nn = 0; nn < K; nn += BLOCK_N)
            {
                size_t Tn = (nn + BLOCK_N <= K) ? BLOCK_N : (K - nn);

                // Within each block, apply 2x2 unrolled kernel if possible.
                size_t i = 0;
                for (; i + 1 < Tm; i += 2)
                {
                    size_t j = 0;
                    for (; j + 1 < Tn; j += 2)
                    {
                        microkernel_2x2(Tk,
                                        K, // using full ldb stride (row-major B)
                                        &A[(mm + i) * N + kk], N,
                                        &B[kk * K + (nn + j)], K,
                                        &C[(mm + i) * K + (nn + j)], K);
                    }
                    // Remaining column (if Tn odd)
                    for (; j < Tn; ++j)
                    {
                        double c0 = 0.0, c1 = 0.0;
                        for (size_t p = 0; p < Tk; ++p)
                        {
                            double a0 = A[(mm + i) * N + (kk + p)];
                            double a1 = A[(mm + i + 1) * N + (kk + p)];
                            double b = B[(kk + p) * K + (nn + j)];
                            c0 += a0 * b;
                            c1 += a1 * b;
                        }
                        C[(mm + i) * K + (nn + j)] += c0;
                        C[(mm + i + 1) * K + (nn + j)] += c1;
                    }
                }
                // Remaining row (if Tm odd)
                for (; i < Tm; ++i)
                {
                    for (size_t j = 0; j < Tn; ++j)
                    {
                        double acc = 0.0;
                        for (size_t p = 0; p < Tk; ++p)
                        {
                            acc += A[(mm + i) * N + (kk + p)] * B[(kk + p) * K + (nn + j)];
                        }
                        C[(mm + i) * K + (nn + j)] += acc;
                    }
                }
            }
        }
    }
}
static int nearly_equal(double a, double b, double rel_tol)
{
    double diff = fabs(a - b);
    double largest = (fabs(a) > fabs(b) ? fabs(a) : fabs(b));
    return diff <= largest * rel_tol + 1e-12; /* small absolute epsilon */
}
int main(int argc, char **argv)
{
    size_t M = 1000, N = 1000, K = 1000;
    if (argc == 4)
    {
        M = (size_t)strtoull(argv[1], NULL, 10);
        N = (size_t)strtoull(argv[2], NULL, 10);
        K = (size_t)strtoull(argv[3], NULL, 10);
    }
    printf("Dimensions: M=%zu N=%zu K=%zu\n", M, N, K);
    srand((unsigned)time(NULL));
    double *A = xmalloc(M * N);
    double *B = xmalloc(N * K);
    double *C_ref = xmalloc(M * K);
    double *C_opt = xmalloc(M * K);
    fill_random(A, M * N);
    fill_random(B, N * K);
    clock_t t0 = clock();
    gemm_naive(M, N, K, A, B, C_ref);
    clock_t t1 = clock();
    ///////////////////////////////////////////
    gemm_blocked(M, N, K, A, B, C_opt); /////////////////////////////////////////
    ///////////////////////////////////////////
    clock_t t2 = clock();
    // Verify a few random entries for correctness.
    int errors = 0;
    for (size_t i = 0; i < M && errors < 10; ++i)
    {
        for (size_t j = 0; j < K && errors < 10; ++j)
        {
            double r = C_ref[i * K + j];
            double o = C_opt[i * K + j];
            if (!nearly_equal(r, o, 1e-9))
            {
                fprintf(stderr, "Mismatch at (%zu,%zu): ref=%g opt=%g\n", i, j, r, o);
                ++errors;
            }
        }
    }

    double naive_ms = 1000.0 * (double)(t1 - t0) / CLOCKS_PER_SEC;
    double opt_ms = 1000.0 * (double)(t2 - t1) / CLOCKS_PER_SEC;
    printf("Naive time: %.3f ms\n", naive_ms);
    printf("Blocked+microkernel time: %.3f ms\n", opt_ms);
    if (errors == 0)
    {
        printf("Verification: PASS\n");
    }
    else
    {
        printf("Verification: FAIL (errors=%d)\n", errors);
    }

    free(A);
    free(B);
    free(C_ref);
    free(C_opt);

    return errors ? EXIT_FAILURE : EXIT_SUCCESS;
}
